import os
import cv2
import numpy as np
from tqdm import tqdm
import random
import argparse

def resize_with_padding(img, target_size=(640, 640)):
    h, w = img.shape[:2]
    scale = min(target_size[0] / h, target_size[1] / w)
    nh, nw = int(h * scale), int(w * scale)
    resized = cv2.resize(img, (nw, nh))

    top = (target_size[0] - nh) // 2
    bottom = target_size[0] - nh - top
    left = (target_size[1] - nw) // 2
    right = target_size[1] - nw - left

    return cv2.copyMakeBorder(resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=[0, 0, 0])

def process_images(input_dir, output_train_dir, output_valid_dir, valid_ratio=0.2, target_size=(640, 640)):
    os.makedirs(output_train_dir, exist_ok=True)
    os.makedirs(output_valid_dir, exist_ok=True)

    all_images = [f for f in os.listdir(input_dir) if f.lower().endswith(('.jpg', '.png', '.jpeg'))]
    random.shuffle(all_images)

    split_idx = int(len(all_images) * (1 - valid_ratio))
    train_imgs = all_images[:split_idx]
    valid_imgs = all_images[split_idx:]

    def save_images(img_list, out_dir):
        for filename in tqdm(img_list, desc=f"Processing -> {out_dir}"):
            img_path = os.path.join(input_dir, filename)
            img = cv2.imread(img_path)
            if img is None:
                print(f"Warning: Cannot read {img_path}")
                continue
            processed_img = resize_with_padding(img, target_size)
            cv2.imwrite(os.path.join(out_dir, filename), processed_img)

    save_images(train_imgs, output_train_dir)
    save_images(valid_imgs, output_valid_dir)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_dir', type=str, required=True, help='Folder with original images')
    parser.add_argument('--output_train_dir', type=str, required=True, help='Where to save training images')
    parser.add_argument('--output_valid_dir', type=str, required=True, help='Where to save validation images')
    parser.add_argument('--size', type=int, default=640, help='Target image size (square)')
    parser.add_argument('--valid_ratio', type=float, default=0.2, help='Validation split ratio')
    args = parser.parse_args()

    process_images(
        input_dir=args.input_dir,
        output_train_dir=args.output_train_dir,
        output_valid_dir=args.output_valid_dir,
        valid_ratio=args.valid_ratio,
        target_size=(args.size, args.size)
    )
